import serial
import struct
import time
import sys
import os
from tqdm import tqdm
import threading
import keyboard
import signal

# Flash地址定义
BOOTLOADER_START_ADDR = 0x08000000
APP_START_ADDR = 0x08004000  # 应用程序起始地址

# OTA命令定义
OTA_CMD_START = 0xA0
OTA_CMD_DATA  = 0xA1
OTA_CMD_END   = 0xA2
OTA_CMD_ABORT = 0xA3
OTA_CMD_RESUME = 0xA5

# OTA响应定义
OTA_RESP_OK    = 0x50
OTA_RESP_ERROR = 0x51
OTA_RESP_READY = 0x52
OTA_RESP_RESUME = 0x53

class OTAUpdater:
    def __init__(self, port, baudrate=115200, chunk_size=1024):
        self.port = port
        self.baudrate = baudrate
        # 确保chunk_size是4的倍数
        self.chunk_size = (chunk_size // 4) * 4
        self.ser = None
        self.is_paused = False
        self.pause_event = threading.Event()
        self.connection_lost = False
        self.auto_reconnect = True

    def connect(self):
        """连接串口"""
        try:
            print(f"Trying to connect to {self.port}...")
            self.ser = serial.Serial(
                port=self.port,
                baudrate=self.baudrate,
                bytesize=serial.EIGHTBITS,
                parity=serial.PARITY_NONE,
                stopbits=serial.STOPBITS_ONE,
                timeout=0.1,           # 读超时0.1秒
                write_timeout=1.0      # 写超时1秒
            )
            print(f"Successfully connected to {self.port} at {self.baudrate} baud")
            return True
        except serial.SerialException as e:
            print(f"Failed to connect to {self.port}: {str(e)}")
            return False
        except Exception as e:
            print(f"Unexpected error while connecting: {str(e)}")
            return False

    def disconnect(self):
        """关闭串口连接"""
        if self.ser and self.ser.is_open:
            self.ser.close()
            print("Serial port closed")

    def calculate_checksum(self, data):
        """计算校验和"""
        return sum(data) & 0xFFFF

    def wait_response(self, expected_resp=OTA_RESP_OK, timeout=0.5):
        """等待并验证响应"""
        print(f"Waiting for response (expected: 0x{expected_resp:02X})...")
        start_time = time.time()
        
        # 清空接收缓冲区
        self.ser.reset_input_buffer()
        
        while (time.time() - start_time) < timeout:
            if self.ser.in_waiting:
                resp = self.ser.read(1)[0]
                print(f"Received response: 0x{resp:02X}")
                if resp == expected_resp:
                    return True
                else:
                    print(f"Error: Unexpected response: 0x{resp:02X}")
                    return False
            time.sleep(0.01)  # 降低CPU占用
        print("Error: Response timeout")
        return False

    def send_packet_with_retry(self, packet, expected_resp=OTA_RESP_OK, timeout=0.5, max_retries=3):
        """发送数据包并处理重试"""
        for retry in range(max_retries):
            try:
                # 检查连接状态
                if not self.check_connection():
                    self.handle_connection_lost()
                    return False
                
                if retry > 0:
                    print(f"\nRetrying... (attempt {retry + 1}/{max_retries})")
                
                self.ser.write(packet)
                if self.wait_response(expected_resp, timeout):
                    return True
                
                if retry < max_retries - 1:
                    time.sleep(0.1 * (retry + 1))
                
            except serial.SerialException as e:
                print(f"串口通信错误: {str(e)}")
                self.handle_connection_lost()
                return False
            except Exception as e:
                print(f"发送数据时出错: {str(e)}")
                if retry < max_retries - 1:
                    time.sleep(0.1 * (retry + 1))
                continue
        
        print(f"在 {max_retries} 次尝试后失败")
        return False

    def calculate_crc32(self, data):
        """计算CRC32校验和"""
        crc = 0xFFFFFFFF
        for b in data:
            crc = crc ^ b
            for _ in range(8):
                if crc & 1:
                    crc = (crc >> 1) ^ 0xEDB88320
                else:
                    crc = crc >> 1
        return ~crc & 0xFFFFFFFF

    def send_start_cmd(self, firmware_data):
        """发送开始升级命令"""
        print("\nStep 1: Sending start command (0xA0)...")
        try:
            # 计算固件CRC32
            firmware_crc = self.calculate_crc32(firmware_data)
            firmware_size = len(firmware_data)
            
            # 准备START命令数据包
            cmd_data = struct.pack('<HHB',  # 小端格式
                0xAA55,                     # sync
                9,                          # length (cmd + size + crc32)
                OTA_CMD_START,              # cmd
            )
            # 添加固件信息
            cmd_data += struct.pack('<II',
                firmware_size,              # 固件大小
                firmware_crc,               # CRC32校验值
            )
            
            # 计算校验和
            checksum = self.calculate_checksum(cmd_data)
            # 添加校验和到数据包末尾
            packet = cmd_data + struct.pack('<H', checksum)
            
            # 发送START命令并等待响应
            print(f"Firmware size: {firmware_size} bytes")
            print(f"Firmware CRC32: 0x{firmware_crc:08X}")
            print("Sending START command...")
            
            self.ser.write(packet)
            
            # 等待响应，可能是0x50(需要重启)或0x52(已就绪)
            start_time = time.time()
            timeout = 5.0  # 设置较长的超时时间
            
            while (time.time() - start_time) < timeout:
                if self.ser.in_waiting:
                    resp = self.ser.read(1)[0]
                    print(f"Received response: 0x{resp:02X}")
                    
                    if resp == OTA_RESP_OK:  # 0x50: 需要重启
                        print("Device will restart to enter bootloader...")
                        time.sleep(3)  # 等待设备重启
                        
                        # 清空接收缓冲区
                        while self.ser.in_waiting:
                            data = self.ser.read(self.ser.in_waiting)
                            print(f"Received during restart: {data.hex()}")
                        
                        # 重新发送START命令（带重试）
                        print("Sending START command again after restart...")
                        if not self.send_packet_with_retry(packet, OTA_RESP_READY, timeout=5.0):
                            print("Error: Device did not respond with READY after restart")
                            return False
                        return True
                        
                    elif resp == OTA_RESP_READY:  # 0x52: 已就绪
                        print("Device is ready for upgrade")
                        return True
                        
                    else:  # 意外的响应
                        print(f"Error: Unexpected response: 0x{resp:02X}")
                        return False
                        
                time.sleep(0.01)  # 降低CPU占用
                
            print("Error: Response timeout")
            return False
            
        except Exception as e:
            print(f"Error sending start command: {str(e)}")
            return False

    def calculate_timeout(self, data_size):
        """根据数据大小和波特率计算超时时间
        Args:
            data_size: 数据包大小(字节)
        Returns:
            timeout: 超时时间(秒)
        """
        # 计算传输一个字节需要的位数 (1起始位 + 8数据位 + 1停止位)
        bits_per_byte = 10
        
        # 计算传输所需时间(秒)
        transfer_time = (data_size * bits_per_byte) / self.baudrate
        
        # Flash写入时间估算 (按每字节1ms计算)
        flash_time = data_size / 1000.0
        
        # 加上一些余量
        timeout = transfer_time + flash_time + 0.5
        
        # 确保最小超时时间
        return max(timeout, 0.5)

    def send_firmware_data(self, firmware_data, resume_offset=0):
        """发送固件数据，支持暂停和断点续传"""
        total_size = len(firmware_data)
        sent_size = resume_offset
        
        # 计算分片相关参数
        FRAGMENT_SIZE = 1024  # 每个分片的大小
        total_fragments = (total_size + FRAGMENT_SIZE - 1) // FRAGMENT_SIZE
        current_fragment = resume_offset // FRAGMENT_SIZE
        
        print(f"Total fragments: {total_fragments}")
        
        with tqdm(total=total_size, initial=resume_offset, 
                  unit='B', unit_scale=True) as pbar:
            while sent_size < total_size:
                # 检查是否暂停或连接断开
                if self.is_paused or self.connection_lost:
                    if self.connection_lost:
                        print("\n等待串口重新连接...")
                    else:
                        print("\n升级已暂停，按 'r' 继续...")
                    self.pause_event.wait()  # 等待唤醒信号
                    print("升级继续进行")
                    continue
                
                # 检查连接状态
                if not self.check_connection():
                    self.handle_connection_lost()
                    continue
                
                # 计算当前分片大小
                chunk_size = min(FRAGMENT_SIZE, total_size - sent_size)
                chunk_data = firmware_data[sent_size:sent_size + chunk_size]
                current_fragment += 1
                
                # 确定分片标志
                if total_fragments == 1:
                    fragment_flag = 0x04  # FRAGMENT_SINGLE
                elif current_fragment == 1:
                    fragment_flag = 0x01  # FRAGMENT_START
                elif current_fragment == total_fragments:
                    fragment_flag = 0x03  # FRAGMENT_END
                else:
                    fragment_flag = 0x02  # FRAGMENT_MIDDLE
                
                # 准备分片头
                fragment_header = struct.pack('<HHHBB',
                    current_fragment,      # fragment_id
                    total_fragments,       # total_fragments
                    chunk_size,            # fragment_size
                    fragment_flag,         # flag
                    0                      # reserved
                )
                
                # 准备数据包
                header = struct.pack('<HHB',  # 小端格式
                    0xAA55,                     # sync
                    len(fragment_header) + chunk_size + 1,  # length
                    OTA_CMD_DATA,               # cmd
                )
                
                # 组装完整数据包
                packet_data = header + fragment_header + chunk_data
                
                # 计算并添加校验和
                checksum = self.calculate_checksum(packet_data)
                packet = packet_data + struct.pack('<H', checksum)
                
                # 发送数据包
                timeout = self.calculate_timeout(len(packet))
                if not self.send_packet_with_retry(packet, OTA_RESP_OK, timeout=timeout):
                    if not self.connection_lost:
                        print("\n发送固件数据失败")
                        return False
                    continue
                
                sent_size += chunk_size
                pbar.update(chunk_size)
                self.save_progress(sent_size)
        
        return True
    
    def save_progress(self, current_size):
        """保存升级进度到文件"""
        try:
            with open('.ota_progress', 'w') as f:
                f.write(str(current_size))
        except Exception as e:
            print(f"Warning: Failed to save progress: {str(e)}")
    
    def load_progress(self):
        """从文件加载上次的升级进度"""
        try:
            if os.path.exists('.ota_progress'):
                with open('.ota_progress', 'r') as f:
                    return int(f.read().strip())
        except Exception:
            pass
        return 0

    def send_end_cmd(self):
        """发送结束命令"""
        print("\nStep 3: Sending end command (0xA2)...")
        try:
            # 准备END命令数据包
            cmd_data = struct.pack('<HHB',  # 小端格式
                0xAA55,                     # sync
                1,                          # length (cmd + addr)
                OTA_CMD_END,                # cmd
            )
            # 计算校验和
            checksum = self.calculate_checksum(cmd_data)
            # 添加校验和到数据包末尾
            packet = cmd_data + struct.pack('<H', checksum)
            
            # 发送END命令（带重试）
            return self.send_packet_with_retry(packet, OTA_RESP_OK)
            
        except Exception as e:
            print(f"Error sending end command: {str(e)}")
            return False

    def check_resume_point(self):
        """检查是否可以断点续传"""
        print("\nChecking resume point...")
        try:
            # 准备RESUME命令数据包
            cmd_data = struct.pack('<HHB',
                0xAA55,         # sync
                1,              # length
                OTA_CMD_RESUME, # cmd
            )
            checksum = self.calculate_checksum(cmd_data)
            packet = cmd_data + struct.pack('<H', checksum)
            
            # 发送RESUME命令
            self.ser.write(packet)
            
            # 等待响应
            start_time = time.time()
            timeout = 1.0
            
            while (time.time() - start_time) < timeout:
                if self.ser.in_waiting >= 5:  # 等待5字节响应
                    resp = self.ser.read(5)
                    if resp[0] == OTA_RESP_RESUME:
                        resume_size = struct.unpack('<I', resp[1:5])[0]
                        print(f"Found resume point at offset: {resume_size}")
                        return resume_size
                    break
                time.sleep(0.01)
                
            return 0
            
        except Exception as e:
            print(f"Error checking resume point: {str(e)}")
            return 0

    def update_firmware(self, firmware_path):
        """执行固件更新流程"""
        try:
            if not os.path.exists(firmware_path):
                print(f"Error: Firmware file not found: {firmware_path}")
                return False
                
            print(f"\nReading firmware file: {firmware_path}")
            with open(firmware_path, 'rb') as f:
                firmware_data = f.read()
            
            print(f"Firmware size: {len(firmware_data)} bytes")
            
            if not self.connect():
                return False
            
            # 检查断点续传
            resume_offset = self.check_resume_point()
            
            # 发送开始命令
            if not self.send_start_cmd(firmware_data):
                return False
            
            # 从断点处发送固件数据
            if not self.send_firmware_data(firmware_data, resume_offset):
                return False
            
            time.sleep(2)
            
            if not self.send_end_cmd():
                return False
            
            print("\nFirmware update completed successfully!")
            return True
            
        except Exception as e:
            print(f"\nUpdate failed: {str(e)}")
            return False
            
        finally:
            self.disconnect()

    def pause_update(self):
        """暂停升级"""
        if not self.is_paused:
            print("\n正在暂停固件升级...")
            self.is_paused = True
            self.pause_event.clear()
            if self.connection_lost:
                print("串口连接断开，等待重新连接...")
            else:
                print("升级已暂停，按 'r' 继续")

    def resume_update(self):
        """继续升级"""
        if self.is_paused and not self.connection_lost:
            print("\n正在继续固件升级...")
            self.is_paused = False
            self.pause_event.set()
            print("升级继续进行")

    def check_connection(self):
        """检查串口连接状态"""
        if not self.ser or not self.ser.is_open:
            return False
        try:
            # 尝试获取串口状态
            self.ser.in_waiting
            return True
        except (serial.SerialException, IOError):
            return False
    
    def auto_reconnect_handler(self):
        """自动重连处理"""
        while self.auto_reconnect and self.connection_lost:
            print("\n尝试重新连接串口...")
            if self.connect():
                print("串口重新连接成功!")
                self.connection_lost = False
                self.resume_update()
                break
            time.sleep(1)  # 等待1秒后重试
    
    def handle_connection_lost(self):
        """处理串口断开"""
        if not self.connection_lost:
            self.connection_lost = True
            print("\n串口连接已断开!")
            self.pause_update()
            
            # 启动重连线程
            reconnect_thread = threading.Thread(target=self.auto_reconnect_handler)
            reconnect_thread.daemon = True
            reconnect_thread.start()

def keyboard_listener(updater):
    """键盘监听线程"""
    try:
        # 尝试导入msvcrt (Windows)
        import msvcrt
        
        while True:
            if msvcrt.kbhit():
                key = msvcrt.getch().decode().lower()
                handle_key(key, updater)
            time.sleep(0.1)
            
    except ImportError:
        # 非Windows系统使用tty
        import tty
        import termios
        import sys
        
        def getch():
            fd = sys.stdin.fileno()
            old_settings = termios.tcgetattr(fd)
            try:
                tty.setraw(sys.stdin.fileno())
                ch = sys.stdin.read(1)
            finally:
                termios.tcsetattr(fd, termios.TCSADRAIN, old_settings)
            return ch
        
        while True:
            key = getch().lower()
            handle_key(key, updater)
            time.sleep(0.1)

def handle_key(key, updater):
    """处理按键"""
    if key == 'p':
        print("\n用户请求暂停...")
        updater.pause_update()
    elif key == 'r':
        print("\n用户请求继续...")
        updater.resume_update()
    elif key == 'q':
        print("\n用户请求退出")
        os._exit(0)

def main():
    # 设置信号处理
    def signal_handler(sig, frame):
        print("\n收到中断信号，正在退出...")
        os._exit(0)
    
    signal.signal(signal.SIGINT, signal_handler)
    
    firmware_path = "build/application.bin"
    print(f"\n开始OTA升级过程...")
    print(f"固件文件: {firmware_path}")
    print("\n控制命令:")
    print("  'p' - 暂停升级")
    print("  'r' - 继续升级")
    print("  'q' - 退出程序")
    
    updater = OTAUpdater('COM6')
    
    # 启动键盘监听线程
    keyboard_thread = threading.Thread(target=keyboard_listener, args=(updater,))
    keyboard_thread.daemon = True
    keyboard_thread.start()
    
    try:
        if updater.update_firmware(firmware_path):
            print("\nOTA升级成功!")
            sys.exit(0)
        else:
            print("\nOTA升级失败!")
            sys.exit(1)
    except KeyboardInterrupt:
        print("\n用户中断升级")
        sys.exit(1)

if __name__ == '__main__':
    main()
